from fastapi import APIRouter
from pydantic import BaseModel
from llm_config.llm_embed_config import Master

llm_load_file = APIRouter()


class File(BaseModel):
    file_name: str = None
    query_similarity_file_info: str = None
    input_info: str


@llm_load_file.post('/retriever_file')
def retriever_info_function(pdf_info: File):
    master = Master()
    text_splitter = master.load_file_name('pdf', pdf_info.file_name)
    retriever_docs_content = master.retriever_file_info(
        text_splitter, pdf_info.query_similarity_file_info)
    print("retriever_docs_content-->", retriever_docs_content)
    return retriever_docs_content


@llm_load_file.post('/load_pdf')
def load_pdf(pdf_info: File):
    master = Master()
    text_splitter = master.load_file_name('pdf', pdf_info.file_name)
    retriever_docs_content = master.retriever_file_info(
        text_splitter, pdf_info.query_similarity_file_info)
    response = master.embed_chain(retriever_docs_content,
                                  pdf_info.query_similarity_file_info)
    return {"response": response["output_text"]}


@llm_load_file.post('/load_txt')
def load_txt(pdf_info: File):
    master = Master()
    text_splitter = master.load_file_name('pdf', pdf_info.file_name)
    retriever_docs_content = master.retriever_file_info(
        text_splitter, pdf_info.query_similarity_file_info)
    response = master.embed_chain(retriever_docs_content,
                                  pdf_info.query_similarity_file_info)
    return {"response": response["output_text"]}


@llm_load_file.post('/llm_memory')
def chain_memory(pdf_info: File):
    master = Master()
    history_info = master.chain_memory(pdf_info.input_info)
    # text_splitter = master.load_file_name('pdf', pdf_info.file_name)
    # retriever_docs_content = master.retriever_file_info(
    #     text_splitter, pdf_info.query_similarity_file_info)
    # response = llm_chain.invoke({
    #     # 问答链（QA Chain）：通常需要 input_documents，因为它需要上下文信息来生成答案
    #     "input_documents": retriever_docs_content,
    #     "context": retriever_docs_content,¬
    #     "question_info": pdf_info.input_info
    # })
    print("response--->", history_info)
    return {"response": history_info}
